import httpx
from fastapi import FastAPI
from fastapi.concurrency import asynccontextmanager


@asynccontextmanager
async def lifespan(app: FastAPI):
    # 创建带连接池的客户端
    limits = httpx.Limits(
        max_connections=20,
        max_keepalive_connections=10,
        keepalive_expiry=60,
    )
    timeout = httpx.Timeout(8)
    client = httpx.AsyncClient(limits=limits, timeout=timeout, follow_redirects=False)

    # 显式绑定到应用状态
    app.state.client = client

    try:
        yield
    finally:
        # 确保关闭客户端
        await client.aclose()


class CacheMiddleware:
    def __init__(self, app):  # 添加构造函数
        self.app = app

    async def __call__(self, scope, receive, send):
        async def modified_send(message):
            if message["type"] == "http.response.start":
                headers = dict(message.get("headers", []))

                # 显式检查 headers 是否为空
                if headers is None:
                    headers = {}

                # 设置 cache-control 头
                headers.setdefault(b"cache-control", b"public, max-age=60")

                # 更新 message 的 headers
                message["headers"] = list(headers.items())

            await send(message)

        await self.app(scope, receive, modified_send)  # 传递处理链
